﻿// Copyright (c) 2012-2013 fancidev
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.

using System;
using System.IO;
using System.Runtime.InteropServices;
using System.IO.Compression;

namespace Util.IO.Compression
{
    /// <summary>
    /// Provides methods to compress and decompress a stream using the LZMA
    /// algorithm.
    /// 
    /// This implementation uses a 13-byte header. The first 5 bytes is the
    /// standard 'props' header. The remaining 8 bytes is the decoded length
    /// stored as an UInt64 in little endian.
    /// </summary>
    /// <remarks>
    /// This class requires LzmaLib.dll.
    /// </remarks>
    public class LzmaStream : Stream
    {
        private Stream baseStream;
        private CompressionMode mode;
        private long length;
        private long position;
        private IntPtr decoder;
        private byte[] cache;
        private int cacheSize;
        private int cacheStart;
        private IntPtr alloc;

        public LzmaStream(Stream baseStream, CompressionMode mode)
        {
            if (baseStream == null)
                throw new ArgumentNullException("baseStream");
            if (mode != CompressionMode.Decompress)
                throw new NotSupportedException();

            this.baseStream = baseStream;
            this.mode = mode;
            this.alloc = LzmaNativeMethods.Lzma_GetDefaultAllocator();
            this.cache = new byte[16000];
            this.cacheSize = 0;
            this.cacheStart = 0;

            this.decoder = LzmaNativeMethods.LzmaDec_Create();
            if (decoder == IntPtr.Zero)
                throw new OutOfMemoryException();

            try
            {
                // Read 5-byte props header.
                byte[] props = new byte[5];
                baseStream.ReadFull(props);
                LzmaErrorCode rc = LzmaNativeMethods.LzmaDec_Allocate(
                    decoder, new ArrayWithOffset(props, 0), 5, this.alloc);
                ThrowLzmaException(rc);

                // Read 8-byte decoded length.
                this.length = baseStream.ReadInt64(Endian.Little);
                if (length < 0)
                    throw new NotSupportedException();
                this.position = 0;

                // Initialize the decoder.
                LzmaNativeMethods.LzmaDec_Init(decoder);
            }
            catch (Exception)
            {
                LzmaNativeMethods.LzmaDec_Destroy(this.decoder);
                throw;
            }
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                if (this.decoder != IntPtr.Zero)
                {
                    LzmaNativeMethods.LzmaDec_Free(this.decoder, IntPtr.Zero);
                    LzmaNativeMethods.LzmaDec_Destroy(this.decoder);
                    this.decoder = IntPtr.Zero;
                }
            }
            base.Dispose(disposing);
        }

        public override bool CanRead
        {
            get { return (this.mode == CompressionMode.Decompress); }
        }

        public override bool CanWrite
        {
            get { return (this.mode == CompressionMode.Compress); }
        }

        public override bool CanSeek
        {
            get { return false; }
        }

        public override int Read(byte[] buffer, int offset, int count)
        {
            if (buffer == null)
                throw new ArgumentNullException("buffer");
            if (offset < 0 || offset > buffer.Length)
                throw new ArgumentOutOfRangeException("offset");
            if (count < 0 || count > buffer.Length - offset)
                throw new ArgumentOutOfRangeException("count");

            // Do not read past stream end.
            if (position >= length)
                return 0;
            count = (int)Math.Min(length - position, (long)count);

            // Loop until we read at least one byte, or encounter EOF.
            while (true)
            {
                if (cacheSize == 0)
                {
                    // Read compressed contents into our intermediate cache.
                    int n = baseStream.Read(cache, 0, cache.Length);
                    if (n < 0 || n > cache.Length)
                        throw new ArgumentException("The underlying stream returned invalid count.");
                    if (n == 0) // EOF
                        throw new EndOfStreamException();

                    cacheStart = 0;
                    cacheSize = n;
                }

                // Decompress some data.
                long destLen = count;
                long srcLen = cacheSize;
                LzmaStatus status = LzmaStatus.LZMA_STATUS_NOT_SPECIFIED;
                LzmaErrorCode rc = LzmaNativeMethods.LzmaDec_DecodeToBuf(
                    decoder,
                    new ArrayWithOffset(buffer, offset), ref destLen,
                    new ArrayWithOffset(cache, cacheStart), ref srcLen,
                    LzmaFinishMode.LZMA_FINISH_ANY, out status);
                ThrowLzmaException(rc);

                if (destLen < 0 || destLen > count)
                    throw new ArgumentException("LzmaLib returned invalid count.");
                if (srcLen < 0 || srcLen > cacheSize)
                    throw new ArgumentException("LzmaLib returned invalid count.");

                this.cacheStart += (int)srcLen;
                this.cacheSize -= (int)srcLen;
                if (destLen > 0)
                {
                    position += destLen;
                    return (int)destLen;
                }

                // If nothing was decompressed, find out why.
                if (status == LzmaStatus.LZMA_STATUS_FINISHED_WITH_MARK ||
                    status == LzmaStatus.LZMA_STATUS_MAYBE_FINISHED_WITHOUT_MARK)
                    throw new EndOfStreamException();
            }
        }

        public override void Write(byte[] buffer, int offset, int count)
        {
            throw new NotImplementedException();
        }

        public override void Flush()
        {
            throw new NotImplementedException();
        }

        public override long Seek(long offset, SeekOrigin origin)
        {
            throw new NotSupportedException();
        }

        public override long Position
        {
            get { return this.position; }
            set { throw new NotSupportedException(); }
        }

        /// <summary>
        /// Gets the uncompressed length of the stream.
        /// </summary>
        public override long Length
        {
            get { return this.length; }
        }

        public override void SetLength(long value)
        {
            throw new NotSupportedException();
        }

        internal static void ThrowLzmaException(LzmaErrorCode rc)
        {
            switch (rc)
            {
                case LzmaErrorCode.SZ_OK:
                    return;
                case LzmaErrorCode.SZ_ERROR_DATA:
                    throw new InvalidDataException();
                case LzmaErrorCode.SZ_ERROR_MEM:
                    throw new OutOfMemoryException();
                case LzmaErrorCode.SZ_ERROR_UNSUPPORTED:
                    throw new NotSupportedException();
                case LzmaErrorCode.SZ_ERROR_PARAM:
                    throw new ArgumentException();
                case LzmaErrorCode.SZ_ERROR_INPUT_EOF:
                case LzmaErrorCode.SZ_ERROR_OUTPUT_EOF:
                    throw new EndOfStreamException();
                case LzmaErrorCode.SZ_ERROR_READ:
                case LzmaErrorCode.SZ_ERROR_WRITE:
                    throw new IOException();
                case LzmaErrorCode.SZ_ERROR_CRC:
                case LzmaErrorCode.SZ_ERROR_PROGRESS:
                case LzmaErrorCode.SZ_ERROR_FAIL:
                case LzmaErrorCode.SZ_ERROR_THREAD:
                case LzmaErrorCode.SZ_ERROR_ARCHIVE:
                case LzmaErrorCode.SZ_ERROR_NO_ARCHIVE:
                default:
                    throw new InvalidOperationException(
                        "Lzma error code: " + rc.ToString());
            }
        }
    }

    internal static class LzmaNativeMethods
    {
        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern IntPtr Lzma_GetDefaultAllocator();

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern IntPtr LzmaDec_Create();

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern void LzmaDec_Destroy(IntPtr decoder);

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern LzmaErrorCode LzmaDec_Allocate(
            IntPtr decoder,
            [In, Out] ArrayWithOffset props,
            int propsSize,
            IntPtr alloc);

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern void LzmaDec_Free(IntPtr decoder, IntPtr alloc);

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern void LzmaDec_Init(IntPtr decoder);

        [DllImport("LzmaLib.dll", ExactSpelling = true, CallingConvention = CallingConvention.Cdecl)]
        public static extern LzmaErrorCode LzmaDec_DecodeToBuf(
            IntPtr decoder,
            [In, Out] ArrayWithOffset dest, // byte *
            ref long destLen,               // size_t
            [In, Out] ArrayWithOffset src,  // byte *
            ref long srcLen,                // size_t
            LzmaFinishMode finishMode,
            out LzmaStatus status);
    }

    internal enum LzmaFinishMode
    {
        /// <summary>Finish at any point.</summary>
        LZMA_FINISH_ANY,

        /// <summary>Block must be finished at the end.</summary>
        LZMA_FINISH_END
    }

    internal enum LzmaStatus
    {
        /// <summary>Use main error code instead.</summary>
        LZMA_STATUS_NOT_SPECIFIED,

        /// <summary>Stream was finished with end mark.</summary>
        LZMA_STATUS_FINISHED_WITH_MARK,

        /// <summary>Stream was not finished.</summary>
        LZMA_STATUS_NOT_FINISHED,

        /// <summary>You must provide more input bytes.</summary>
        LZMA_STATUS_NEEDS_MORE_INPUT,

        /// <summary>
        /// There is probability that stream was finished without end mark.
        /// </summary>
        LZMA_STATUS_MAYBE_FINISHED_WITHOUT_MARK  
    }

    internal enum LzmaErrorCode
    {
        SZ_OK = 0,
        SZ_ERROR_DATA = 1,
        SZ_ERROR_MEM = 2,
        SZ_ERROR_CRC = 3,
        SZ_ERROR_UNSUPPORTED = 4,
        SZ_ERROR_PARAM = 5,
        SZ_ERROR_INPUT_EOF = 6,
        SZ_ERROR_OUTPUT_EOF = 7,
        SZ_ERROR_READ = 8,
        SZ_ERROR_WRITE = 9,
        SZ_ERROR_PROGRESS = 10,
        SZ_ERROR_FAIL = 11,
        SZ_ERROR_THREAD = 12,
        SZ_ERROR_ARCHIVE = 16,
        SZ_ERROR_NO_ARCHIVE = 17,
    }

#if false
        public static extern LzmaErrorCode Lzma86CS_GetUnpackSize(
            [MarshalAs(UnmanagedType.LPArray)] byte[] pInputBuffer,
            int iInputStartIndex,
            int iInputLength,
            out int pUnpackSize);

        [DllImport("LzmaLib.dll")]
        private static extern LzmaErrorCode Lzma86CS_Decode(
            [MarshalAs(UnmanagedType.LPArray)] byte[] pOutputBuffer,
            int iOutputStartIndex,
            int nOutputLength,
            [MarshalAs(UnmanagedType.LPArray)] byte[] pInputBuffer,
            int iInputStartIndex,
            int nInputLength);

        public static int Decode(
            byte[] inputBuffer, int inputStartIndex, int inputLength,
            byte[] outputBuffer, int outputStartIndex, int outputLength)
        {
            // Get the unpacked size.
            int unpackSize;
            LzmaErrorCode result = Lzma86CS_GetUnpackSize(inputBuffer, inputStartIndex, inputLength, out unpackSize);
            if (result != LzmaErrorCode.SZ_OK)
                throw new LzmaException(result);

            // If the unpacked size is greater than available buffer, raise an error.
            if (unpackSize > outputLength)
                throw new LzmaException(LzmaErrorCode.SZ_ERROR_OUTPUT_EOF);

            // Do the actual unpacking.
            result = Lzma86CS_Decode(
                outputBuffer, outputStartIndex, unpackSize,
                inputBuffer, inputStartIndex, inputLength);
            if (result != LzmaErrorCode.SZ_OK)
                throw new LzmaException(result);

            // Return the unpacked size.
            return unpackSize;
        }
#endif

#if false
    class LzssDecoder
    {
        [DllImport("LzmaLib.dll")]
        private static extern int LzssCS_Decode(
            [MarshalAs(UnmanagedType.LPArray)] byte[] pOutputBuffer,
            int iOutputStartIndex,
            int nOutputLength,
            [MarshalAs(UnmanagedType.LPArray)] byte[] pInputBuffer,
            int iInputStartIndex,
            int nInputLength);

        public static int Decode(
            byte[] inputBuffer, int inputStartIndex, int inputLength,
            byte[] outputBuffer, int outputStartIndex, int outputLength)
        {
            int result = LzssCS_Decode(
                outputBuffer, outputStartIndex, outputLength,
                inputBuffer, inputStartIndex, inputLength);
            if (result < 0)
                throw new InvalidOperationException();

            // Return the unpacked size.
            return result;
        }
    }
#endif
}
